//
//  KDTree.cpp
//  RayTracer
//
//  Created by 俞云康 on 4/14/16.
//  Copyright © 2016 yushroom. All rights reserved.
//

#include "KDTree.hpp"

typedef float Float;

struct KDTreeNode {
    // KDTreeNode Methods
    void InitLeaf(int *primNums, int np, std::vector<int> *primitiveIndices) {
        flags = 3;
        nPrims |= (np << 2);
        // store primitive ids for leaf node
        if (np == 0)
            onePrimitive = 0;
        else if (np == 1)
            onePrimitive = primNums[0];
        else {
            primitiveIndicesOffset = (int)primitiveIndices->size();
            for (int i = 0; i < np; ++i)
                primitiveIndices->push_back(primNums[i]);
        }
    }
    void InitInterior(int axis, int ac, Float s) {
        split = s;
        flags = axis;
        aboveChild |= (ac << 2);
    }
    Float SplitPos() const { return split; }
    int nPrimitives() const { return nPrims >> 2; }
    int SplitAxis() const { return flags & 3; }
    bool IsLeaf() const { return (flags & 3) == 3; }
    int AboveChild() const { return aboveChild >> 2; }
    union {
        Float split;                 // Interior
        int onePrimitive;            // Leaf
        int primitiveIndicesOffset;  // Leaf
    };
    
private:
    union {
        int flags;       // Both
        int nPrims;      // Leaf
        int aboveChild;  // Interior
    };
};

struct KdToDo {
    const KDTreeNode *node;
    float tMin, tMax;
};

enum class EdgeType { Start, End };

struct BoundEdge {
    // BoundEdge Public Methods
    BoundEdge() {}
    BoundEdge(Float t, int primNum, bool starting) : t(t), primNum(primNum) {
        type = starting ? EdgeType::Start : EdgeType::End;
    }
    Float t;
    int primNum;
    EdgeType type;
};

KDTree::KDTree(const vector<shared_ptr<GeometricPrimitive>> &p,
               int isectCost,
               int traversalCost,
               float emptyBonus,
               int maxPrims,
               int maxDepth)
:   isectCost(isectCost),
    traversalCost(traversalCost),
    maxPrims(maxPrims),
    emptyBonus(emptyBonus),
    primitives(p)
{
    nextFreeNode = nAllocedNodes = 0;
    if (maxDepth <= 0)
        maxDepth = (int)std::round(8+1.3f*Log2Int((uint32_t)primitives.size()));
    
    vector<AABB> primBounds;
    primBounds.reserve(primitives.size());
    for (auto& prim : primitives) {
        AABB b = prim->worldBounds();
        bounds.expandBy(b);
        primBounds.push_back(b);
    }
    
    unique_ptr<BoundEdge[]> edges[3];
    for (int i = 0; i < 3; ++i)
        edges[i].reset(new BoundEdge[2 * primitives.size()]);
    unique_ptr<int[]> prims0(new int[primitives.size()]);
    unique_ptr<int[]> prims1(new int[(maxDepth + 1) * primitives.size()]);
    
    unique_ptr<int[]> primNums(new int[primitives.size()]);
    for (size_t i = 0; i < primitives.size(); ++i)
        primNums[i] = i;
    
    buildTree(0, bounds, primBounds, primNums.get(),
              (int)primitives.size(), maxDepth, edges,
              prims0.get(), prims1.get());
}

KDTree::~KDTree() {
    delete [] nodes;
}

bool KDTree::intersect(const Ray& ray, HitInfo * isect) const {
    float tMin = 0, tMax = 0;
    if (!bounds.intersect(ray, &tMin, &tMax))
        return false;
    
    Direction invDir;
	//Direction invDir(1.f / ray.dir.x, 1.f / ray.dir.y, 1 / ray.dir.z);
	invDir.x = 1.f / ray.dir.x;
	invDir.y = 1.f / ray.dir.y;
	invDir.z = 1.f / ray.dir.z;
    constexpr int maxTodo = 64;
    KdToDo todo[maxTodo];
    int todoPos = 0;
    
    // Traverse kd-tree nodes in order for ray
    bool hit = false;
    const KDTreeNode *node = &nodes[0];
    while (node != nullptr) {
        // Bail out if we found a hit closer than the current node
        if (ray.maxt < tMin) break;
        if (!node->IsLeaf()) {
            // Process kd-tree interior node
           
            // Compute parametric distance along ray to split plane
            int axis = node->SplitAxis();
            float tPlane = (node->SplitPos() - ray.o[axis]) * invDir[axis];
            
            // Get node children pointers for ray
            const KDTreeNode *firstChild, *secondChild;
            int belowFirst =
            (ray.o[axis] < node->SplitPos()) ||
            (ray.o[axis] == node->SplitPos() && ray.dir[axis] <= 0);
            if (belowFirst) {
                firstChild = node + 1;
                secondChild = &nodes[node->AboveChild()];
            } else {
                firstChild = &nodes[node->AboveChild()];
                secondChild = node + 1;
            }
            
            // Advance to next child node, possibly enqueue other child
            if (tPlane > tMax || tPlane <= 0)
                node = firstChild;
            else if (tPlane < tMin)
                node = secondChild;
            else {
                // Enqueue _secondChild_ in todo list
                todo[todoPos].node = secondChild;
                todo[todoPos].tMin = tPlane;
                todo[todoPos].tMax = tMax;
                ++todoPos;
                node = firstChild;
                tMax = tPlane;
            }
        } else {
            // Check for intersections inside leaf node
            int nPrimitives = node->nPrimitives();
            if (nPrimitives == 1) {
                const auto& p = primitives[node->onePrimitive];
                if (p->intersect(ray, isect))
                    hit = true;
            } else {
                for (int i = 0; i < nPrimitives; ++i) {
					//Assert(node->primitiveIndicesOffset + i < primitiveIndices.size());
                    int index = primitiveIndices[node->primitiveIndicesOffset + i];
					//Assert(index < primitives.size());
                    const auto& p = primitives[index];
                    // Check one primitive inside leaf node
                    if (p->intersect(ray, isect))
                        hit = true;
                }
            }
            
            // Grab next node to process from todo list
            if (todoPos <= 0) break;
            --todoPos;
            node = todo[todoPos].node;
            tMin = todo[todoPos].tMin;
            tMax = todo[todoPos].tMax;
        }
    }
    return hit;
}

bool KDTree::intersect(const Ray& ray) const {
    float tMin = 0, tMax = 0.0;
    if (!bounds.intersect(ray, &tMin, &tMax)) {
        return false;
    }
    
    Direction invDir(1.f/ray.dir.x, 1.f/ray.dir.y, 1.f/ray.dir.z);
    constexpr int maxTodo = 64;
    KdToDo todo[maxTodo];
    int todoPos = 0;
    const KDTreeNode *node = &nodes[0];
    while (node != nullptr) {
        if (node->IsLeaf()) {
            // Check for shadow ray intersections inside leaf node
            int nPrimitives = node->nPrimitives();
            if (nPrimitives == 1) {
                const auto& p = primitives[node->onePrimitive];
                if (p->intersect(ray)) return true;
            } else {
                for (int i = 0; i < nPrimitives; ++i) {
                    int index = primitiveIndices[node->primitiveIndicesOffset + i];
                    const auto& p = primitives[index];
                    if (p->intersect(ray)) return true;
                }
            }
            
            if (todoPos <= 0) break;
            --todoPos;
            node = todo[todoPos].node;
            tMin = todo[todoPos].tMin;
            tMax = todo[todoPos].tMax;
        } else {
            // Process kd-tree interior node
            
            // Compute parametric distance along ray to split plane
            int axis = node->SplitAxis();
            float tPlane = (node->SplitPos()-ray.o[axis]) * invDir[axis];
            
            // Get node children pointers for ray
            const KDTreeNode *firstChild, *secondChild;
            int belowFirst = (ray.o[axis] < node->SplitPos()) ||
                            (ray.o[axis] == node->SplitPos() && ray.dir[axis] <= 0);
            if (belowFirst) {
                firstChild = node + 1;
                secondChild = nodes + node->AboveChild();
            } else {
                firstChild = nodes + node->AboveChild();
                secondChild = node + 1;
            }
            
            // Advance to next child node, possibly enqueue other child
            if (tPlane > tMax || tPlane <= 0)
                node = firstChild;
            else if (tPlane < tMin)
                node = secondChild;
            else {
                todo[todoPos].node = secondChild;
                todo[todoPos].tMin = tPlane;
                todo[todoPos].tMax = tMax;
                ++todoPos;
                node = firstChild;
                tMax = tPlane;
            }
            if (node - nodes > nAllocedNodes) {
                info("here\n");
                //node = nullptr;
            }
        }
    }
    return false;
}


void KDTree::buildTree(int nodeNum,
                       const AABB& nodeBounds,
                       const vector<AABB>& allPrimBounds,
                       int *primNums,
                       int nPrimitives,
                       int depth,
                       const unique_ptr<BoundEdge[]> edges[3],
                       int *prims0,
                       int *prims1,
                       int badRefines)
{
    Assert(nodeNum == nextFreeNode);
    if (nextFreeNode == nAllocedNodes) {
        int nNewAllocNodes = std::max(2*nAllocedNodes, 512);
        KDTreeNode *n = new KDTreeNode[nNewAllocNodes];
        if (nAllocedNodes > 0) {
            memcpy(n, nodes, nAllocedNodes * sizeof(KDTreeNode));
            delete [] nodes;
        }
        nodes = n;
        nAllocedNodes = nNewAllocNodes;
    }
    ++nextFreeNode;
    
    // Initialize leaf node if termination criteria met
    if (nPrimitives <= maxPrims || depth == 0) {
        nodes[nodeNum].InitLeaf(primNums, nPrimitives, &primitiveIndices);
        return;
    }
    
    // Initialize interior node and continue recursion
    
    // Choose split axis position for interior node
    int bestAxis = -1, bestOffset = -1;
    float bestCost = INFINITY;
    float oldCost = isectCost * float(nPrimitives);
    float totalSA = nodeBounds.surfaceArea();
    float invTotalSA = 1.f / totalSA;
    Direction d = nodeBounds.diagonal();
    
    // Choose which axis to split along
    int axis = nodeBounds.maximumExtent();
    int retries = 0;
retrySplit:
    // Initialize edges for _axis_
    for (int i = 0; i < nPrimitives; ++i) {
        int pn = primNums[i];
        const auto& bounds = allPrimBounds[pn];
        edges[axis][2*i]   = BoundEdge(bounds.pmin[axis], pn, true);
        edges[axis][2*i+1] = BoundEdge(bounds.pmax[axis], pn, false);
    }
    
    std::sort(&edges[axis][0], &edges[axis][2*nPrimitives],
              [](const BoundEdge &e0, const BoundEdge &e1) -> bool {
                  return e0.t == e1.t ? int(e0.type) < int(e1.type) : e0.t < e1.t;
              });
    
    // Compute cost of all splits for _axis_ to find best
    int nBelow = 0, nAbove = nPrimitives;
    for (int i = 0; i < 2*nPrimitives; ++i) {
        if (edges[axis][i].type == EdgeType::End)
            --nAbove;
        float edgeT = edges[axis][i].t;
        if (edgeT > nodeBounds.pmin[axis] && edgeT < nodeBounds.pmax[axis]) {
            // Compute cost for split at _i_th edge
            
            // Compute child surface areas for split at _edgeT_
            int otherAxis0 = (axis+1)%3, otherAxis1 = (axis+2)%3;
            float belowSA = 2*(d[otherAxis0]*d[otherAxis1] +
                               (edgeT - nodeBounds.pmin[axis]) * (d[otherAxis0]+d[otherAxis1]));
            float aboveSA = 2*(d[otherAxis0]*d[otherAxis1] +
                               (nodeBounds.pmax[axis]-edgeT) * (d[otherAxis0]+d[otherAxis1]));
            float pBelow = belowSA * invTotalSA;
            float pAbove = aboveSA * invTotalSA;
            float eb = (nAbove == 0 || nBelow == 0) ? emptyBonus : 0;
            float cost = traversalCost + isectCost*(1-eb)*(pBelow*nBelow + pAbove*nAbove);
            
            // Update best split if this is lowest cost so far
            if (cost < bestCost) {
                bestCost = cost;
                bestAxis = axis;
                bestOffset = i;
            }
        }
        if (edges[axis][i].type == EdgeType::Start)
            ++nBelow;
    }
    Assert(nBelow == nPrimitives && nAbove == 0);
    
    // Create leaf if no good splits were found
    if (bestAxis == -1 && retries < 2) {
        ++retries;
        axis = (axis + 1) % 3;
        goto retrySplit;
    }
    if (bestCost > oldCost)
        ++badRefines;
    if ((bestCost > 4*oldCost && nPrimitives < 16) || bestAxis == -1 || badRefines == 3) {
        nodes[nodeNum].InitLeaf(primNums, nPrimitives, &primitiveIndices);
        return;
    }
    
    // Classify primitives with respect to split
    int n0 = 0, n1 = 0;
    for (int i = 0; i < bestOffset; ++i) {
        if (edges[bestAxis][i].type == EdgeType::Start)
            prims0[n0++] = edges[bestAxis][i].primNum;
    }
    for (int i = bestOffset+1; i < 2 * nPrimitives; ++i) {
        if (edges[bestAxis][i].type == EdgeType::End)
            prims1[n1++] = edges[bestAxis][i].primNum;
    }
    
    // Recursively initialize children nodes
    float tSplit = edges[bestAxis][bestOffset].t;
    AABB bounds0 = nodeBounds, bounds1 = nodeBounds;
    bounds0.pmax[bestAxis] = bounds1.pmin[bestAxis] = tSplit;
    buildTree(nodeNum+1, bounds0, allPrimBounds, prims0, n0, depth-1, edges, prims0, prims1+nPrimitives, badRefines);
    int aboveChild = nextFreeNode;
    nodes[nodeNum].InitInterior(bestAxis, aboveChild, tSplit);
    buildTree(aboveChild, bounds1, allPrimBounds, prims1, n1, depth-1, edges, prims0, prims1+nPrimitives, badRefines);
}
